import os
import argparse
from pathlib import Path
import re
from bamstat.coverage import calculate_transcript_coverages, calculate_bed_coverages
from bamstat.data import read_genepred, read_bed
from bamstat.tview import tview_bam
from bamstat.plot import plot_snv


def run_coverage_cmd(args):
    region_format = args.region_format
    if region_format == 'gpe':
        gpes = read_genepred(gpe_file=args.region)
        coverages = calculate_transcript_coverages(bam_file=args.bam, gpes=gpes, thread=args.thread)
    else:
        beds = read_bed(bed_file=args.region)
        coverages = calculate_bed_coverages(bam_file=args.bam, beds=beds, thread=args.thread)
    outfile = args.out if args.out.endswith('.gz') else f'{args.out}.gz'
    coverages.to_csv(outfile, sep='\t', index=False, compression="gzip")


def run_plot_cmd(args):
    match = re.match(r'^(\w+):(\d+)-?(\d+)?$', args.pos)
    if not match:
        raise Exception(f'Error: unknown position format: {args.pos}')
    chrom, start, end = match.groups()
    tview = tview_bam(bam=args.bam, ref=args.ref, chrom=chrom, start=int(start), end=int(end or start), extend=args.extend, depth=args.depth)
    plot_snv(tview=tview, extend=args.extend, img=args.out)


def get_options():
    parser = argparse.ArgumentParser(description='BAM Stat')
    sub_parsers = parser.add_subparsers()
    cov_parser = sub_parsers.add_parser('cov', help='BAM Coverage')
    cov_parser.add_argument('--bam', '-i', required=True, help='input, BAM file')
    cov_parser.add_argument('--region', '-r', required=True, help='input, region file, BED or GenePred')
    cov_parser.add_argument('--region_format', '-f', choices=('gpe', 'bed'), required=True, help='parameter, region file format')
    cov_parser.add_argument('--out', '-o', required=True, help='output, Output file')
    cov_parser.add_argument('--thread', '-t', type=int, default=4, help='parameter, Thread number')
    cov_parser.set_defaults(func=run_coverage_cmd)
    plot_parser = sub_parsers.add_parser('plot', help='Plot BAM')
    plot_parser.add_argument('--bam', '-i', required=True, help='input, BAM file')
    plot_parser.add_argument('--ref', '-r', required=True, help='input, Reference genome file')
    plot_parser.add_argument('--out', '-o', required=True, help='output, Output image file')
    plot_parser.add_argument('--pos', '-p', required=True, help='parameter, position info, chrom:start or chrom:start-end')
    plot_parser.add_argument('--extend', '-e', type=int, default=50, help='input, BAM file')
    plot_parser.add_argument('--depth', '-d', type=int, help='input, BAM file')
    plot_parser.set_defaults(func=run_plot_cmd)
    return parser.parse_args()


def main():
    args = get_options()
    args.func(args)


if __name__ == '__main__':
    main()
